
def parse_example(example_path):
	with open(example_path) as fp:
		example = fp.read()
	index_problem = example.find('**Problem**:')
	assert index_problem >= 0, index_problem
	index_solution = example.find('**Solution**:')
	assert index_solution > index_problem
	index_reasoning = example.find('**Reasoning**:')
	assert index_reasoning > index_solution
	index_insight = example.find('**Insight**:')
	assert index_insight > index_reasoning
	problem = example[index_problem:index_solution]
	solution = example[index_solution:index_reasoning]
	reasoning = example[index_reasoning:index_insight]
	insight = example[index_insight:]
	return [problem, solution, reasoning, insight]


from llm import gen
from util import color, parallel
from exp import data, run
import random, copy, uuid, json, os


# First use prompt "List the major reasoning steps according to solution for the problem." extract reasoning steps from the ground-truth solution of the problem
def step11(model_id, examples, sample, verbal=True):
	prompt_templ = '''#TASK DESCRIPTION: 
List the major reasoning steps in the solution for the problem.

## EXAMPLE 1

{example_1}

## EXAMPLE 2

{example_2}

## EXAMPLE 3

{example_3}

## YOUR TASK: Decompose the **Solution** for the **Problem** into 2-4 major reasoning steps in **Reasoning** according to the style and format in the examples.

**Problem**: {problem}

**Solution**: {solution}'''
	random.shuffle(examples)
	prompt = prompt_templ.format(example_1=''.join(examples[0][:3]),
								example_2=''.join(examples[1][:3]),
								example_3=''.join(examples[2][:3]),
								problem=sample['problem'],
								solution=sample['cot'],
								)
	if verbal:
		color.text(prompt, 'yellow')
	messages = gen.make_prompt(user_prompt=prompt)
	response, _, usage = gen.generate(model_id, messages, enable_thinking=False, verbal=verbal)
	return response, usage


# Then use prompt "Explain the insight behind how to conceive one or more of your consecutive reasoning steps." to generate Insights for the reasoning steps.
def step12(model_id, examples, sample, reasoning_steps, verbal=True):
	prompt_templ = '''#TASK DESCRIPTION: 
Explain the insights behind how to conceive your consecutive reasoning steps.

## EXAMPLE 1

{example_1}

## EXAMPLE 2

{example_2}

## EXAMPLE 3

{example_3}

## YOUR TASK: Generate the **Insights** underlying the **Reasoning** steps according to the style and format in the examples. The insights generated should be generally useful to guide reasoning in future problems with different entities and values.

**Problem**: {problem}

{reasoning_steps}'''
	random.shuffle(examples)
	prompt = prompt_templ.format(example_1=''.join(examples[0][:1] + examples[0][2:4]),
								example_2=''.join(examples[1][:1] + examples[0][2:4]),
								example_3=''.join(examples[2][:1] + examples[0][2:4]),
								problem=sample['problem'],
								reasoning_steps=reasoning_steps,
								)
	if verbal:
		color.text(prompt, 'yellow')
	messages = gen.make_prompt(user_prompt=prompt)
	response, _, usage = gen.generate(model_id, messages, enable_thinking=False, verbal=verbal)
	return response, usage


def parse_reasoning_steps(steps, prefix='- Step'):
	return [s.split(':', 1)[1] for s in steps.split(prefix)[1:]]

def parse_insights(insights, prefix='- For Reasoning Step'):
	def parse_1(s):
		s = [s1.split(':', 1) for s1 in s.split('\n  - ')[1:]]
		s = {k:v for k,v in s if k in ('Situation','Goal')}
		assert 'Situation' in s and 'Goal' in s, s
		return s
	return [parse_1(s) for s in insights.split(prefix)[1:]]

def seed_examples():
	example1 = parse_example('com/init-examples/example1.md')
	example2 = parse_example('com/init-examples/example2.md')
	example3 = parse_example('com/init-examples/example3.md')
	return [example1, example2, example3]

# - Step 1 (create.py): build an initial library of Insights
def step1(test_set='MATH500', model_id=None, repeat=10, verbal=False, examples_path='../OUT/created-examples'):
	if model_id is None:
		model_id = run.model_ids[0]
	if not os.path.isdir(examples_path):
		os.mkdir(examples_path)
	outfile = f'{examples_path}/step1_{test_set}.jsonl'
	if os.path.isfile(outfile):
		return
	examples = seed_examples()
	samples = data.get_data(test_set)
	def step1_task(sample, retries=10):
		for i in range(retries):
			try:
				reasoning_steps, usage1 = step11(model_id, examples, sample, verbal=verbal)
				insights, usage2 = step12(model_id, examples, sample, reasoning_steps, verbal=verbal)
				reasoning_steps = parse_reasoning_steps(reasoning_steps)
				insights = parse_insights(insights)
				assert len(reasoning_steps) == len(insights)
				sample = copy.deepcopy(sample)
				sample['reasoning_steps'] = reasoning_steps
				sample['insights'] = insights
				sample['uuid'] = str(uuid.uuid1())
				return sample
			except:
				pass
	if repeat > 1:
		samples = samples * 10
	res = parallel.run(step1_task, samples, num_tasks=50)
	res = [r for r in res if r is not None]
	res.sort(key=lambda r: r['pid'])
	with open(outfile, 'w') as fp:
		for r in res:
			fp.write(json.dumps(r, ensure_ascii=False) + '\n')

if __name__ == '__main__':
	step1()